import gensim
from gensim.models import Word2Vec
from gensim.utils import simple_preprocess

# 读取语料库
with open('train-data.txt', 'r', encoding='utf-8') as file:
    sentences = [line.strip() for line in file]

# 预处理：分词并去除标点符号
processed_sentences = [simple_preprocess(sentence) for sentence in sentences]

# 打印预处理后的句子
for sentence in processed_sentences:
    print(sentence)

# 训练 Word2Vec 模型（使用 Skip-gram）
model = Word2Vec(
    sentences=processed_sentences,  # 输入数据
    vector_size=100,  # 嵌入维度
    window=5,  # 上下文窗口大小
    min_count=1,  # 忽略出现频率低于此值的词汇
    workers=4,  # 使用多线程加速训练
    sg=1  # 使用 Skip-gram 模型 (sg=0 表示使用 CBOW)
)

# 保存模型
model.save("word2vec.model")

# 加载预训练的 Word2Vec 模型
model = Word2Vec.load("word2vec.model")

# 查找与 "python" 最相似的词汇
similar_words = model.wv.most_similar("python", topn=5)
print(similar_words)

# 计算 "python" 和 "language" 之间的相似度
similarity = model.wv.similarity("python", "language")
print(f"'python'和'language'相似度: {similarity:.4f}")

# 获取 "python" 的嵌入向量
vector = model.wv["python"]
print(vector)
